/*
 * Copyright 2021-2022 Open Kunlun Technology <https://www.openkunlun.io>
 */

package io.openkunlun.scaladsl.context

import akka.actor.{ ExtendedActorSystem, Extension, ExtensionId, ExtensionIdProvider }
import akka.http.scaladsl.model.HttpRequest
import io.openkunlun.scaladsl.context.Context.ContextTag

/**
 * @author ericxin.
 */
object AkkaHttpContext extends ExtensionId[AkkaHttpContext] with ExtensionIdProvider {
  override def lookup: ExtensionId[AkkaHttpContext] = AkkaHttpContext
  override def createExtension(system: ExtendedActorSystem) = new AkkaHttpContext(system)
}
class AkkaHttpContext(system: ExtendedActorSystem) extends Extension {

  private val contextPropagation: ContextStorage = ContextStorage.current(system)

  def init(): Unit = {
    contextPropagation.init()
  }

  def store[T](request: HttpRequest)(f: => T): T = {
    val tags = request.headers.map(it => it.name() -> it.value())
      .map(it => ContextTag(it._1, it._2))
    contextPropagation.store(new Context(tags))(f)
  }

  def get(): Context = {
    contextPropagation.get()
  }
}
